#!/usr/local/bin/python

import argparse
import os
import sys
import hashlib
import socket
import datetime
import psycopg2
import pysftp

parser = argparse.ArgumentParser(description= """

DESCRIPTION:
                                 
    Insert into table project_files in postgres or checks for non-existing files.
                                     
    Sub-directories are ignored unless -r option is given
    
--vacuum:
    This option reads table project_files and checks whether a file exists.
    NOTE: It will only check that a file with that name and location exists, it
    will not check the md5sum! (Yet).
    Only files with project_files.removed=False are checked (at the moment)
    
EXAMPLES:

    - Assign file1.txt and file2.txt and files in dir Tritume to myproject
    (but don't go into subdirs)
    $ project_files.py -p 'myproject' -f file1.txt file2.txt
     
    - Get the list of files from stdin (files in current dir with .txt ext):
    $ ls *.txt | project_files.py -p 'myproject' -f -

TODO:

    - Vacuum: Returns: 
    Exception AttributeError: "'Connection' object has no attribute '_tranport_live'" in <bound method Connection.__del__ of <pysftp.Connection object at 0x1029250>> ignored
    - Vacumm: Add support for project specific vacuuming?
    
""", formatter_class= argparse.RawTextHelpFormatter)

parser.add_argument('--vacuum',
                   action= 'store_true',
                   help='''Read table project_files and set remved=False the files that could not be found.
                   
                   ''')
parser.add_argument('--connection', '-c',
                   type= str,
                   required= False,
                   help='''A string to be passed to psycopg2.connect() to connect to the database. E.g. "dbname=testdb user=me password= mypwd host= localhost"
If None, connection parameters will be read from ~/.psycopgpass

''')
parser.add_argument('--project', '-p',
                   type= str,
                   required= False,
                   help='''The project to which the files should be assigned.
                   
                   ''')
parser.add_argument('--description', '-d',
                   type= str,
                   required= False,
                   help='''A description for this file(s)
                   
                   ''')
parser.add_argument('--recursive', '-r',
                   action='store_true',
                   help= '''Recurse into the directories found in the list of files? Default false
                   
                   ''')
parser.add_argument('--hidden', '-H',
                    action='store_true',
                    help= """Include hidden files and dirs. (Recognized by leading '.' so not platform independent.)
                    
                    """)
parser.add_argument('-f', '--files',
                   type= str, ## argparse.FileType('r'),
                   nargs= '+',
                   required= False,
                   help= '''Files or directories to include, possibly recursing into the listed dirs. use "--files -" to read the list of files from stdin.
                   
                   ''')
parser.add_argument('-v', '--verbose',
                   action='store_true',
                   help= '''Print to stdout the sql commands sent to postgres
                   
                   ''')
args = parser.parse_args()

" ------------------------ [ checks on arguments ]---------------------------- "

if args.connection is None:
    try:
        conn= open(os.path.join(os.getenv("HOME"), '.psycopgpass'))
    except:
        sys.exit('I cannot read file %s and -c/--connection parameter is None.\nI cannot connect to any database.' %(os.path.join(os.getenv("HOME"), '.psycopgpass')))
    conn= conn.readlines()
    conn_args= [x.strip() for x in conn if x.strip() != '' and x.strip().startswith('#') is False][0]
    print('\nNOTICE: Connecting to database using parameters in ~/.psycopgpass\n')
else:
    conn_args= args.connection

if args.vacuum is False:
    if args.project is None:
        sys.exit('Argument --project/-p is required.')
    if args.description is None:
        sys.exit('Argument --description/-d is required.')
    if args.files is None:
        sys.exit('Argument --files/-d is required.')
else:
    """
    ____________________________________________________________________________
                             
                             --vacuum MODULE
                             
    Check for non existiting files in table project_files
    ____________________________________________________________________________
    """
    ## sys.exit("The vacuum module is not available yet.")
    
    conn = psycopg2.connect(conn_args)
    cur= conn.cursor()
    cur.execute(""" SELECT DISTINCT hostname, array_agg(os_path_join(path, filename)) FROM project_files WHERE removed = False GROUP BY hostname ORDER BY hostname LIMIT 10 """)
    servers_files= cur.fetchall()
    for host, files in servers_files:
        try:
            srv= pysftp.Connection(host)
        except:
            print("WARNING: Skipping server (I cannot connect): %s" %(host))
            continue
        files_quoted= []
        print('Checking files on server: %s (%s files to check)' %(host, len(files)))
        for f in files:
            " Quote files before passing to bash "
            if '"' in f:
                " Single quotes if filename contains double quotes "
                files_quoted.append("'" + f + "'")
            else:
                " Double quotes all other files "
                files_quoted.append('"' + f + '"')
        ## Test for file existance: Returns a list containing 'exists' or 'removed'
        tester= "for f in %s; do if [ -f $f ]; then echo 'exists'; else echo 'removed'; fi done" %(' '.join(files_quoted))
        ftest= srv.execute(tester)
        ftest= [x.strip() for x in ftest]
        srv.close()
        if len(ftest) != len(files):
            sys.exit('Different number of files tested and returned')
        for test, filepath in zip(ftest, files):
            " Update project_files where file does not exist anymore "
            if test == 'exists':
                continue
            path, filename= os.path.split(filepath)
            update_sql= cur.mogrify('UPDATE project_files SET removed = True WHERE hostname = %s AND filename= %s AND path = %s', (host, filename, path) )
            cur.execute(update_sql)
            if args.verbose:
                print(update_sql)
    cur.close()
    conn.commit()
    conn.close()
    sys.exit()       
""" ______________________ END OF --vacuum MODULE ______________________________

----------------------- Classes and Functions ----------------------------------
"""

def sumfile(fobj):
    '''Returns an md5 hash for an object with read() method.'''
    m = hashlib.md5()
    while True:
        d = fobj.read(8096)
        if not d:
            break
        m.update(d)
    return m.hexdigest()


def md5sum(fname):
    '''Returns an md5 hash for file fname, or stdin if fname is "-".'''
    if fname == '-':
        ret = sumfile(sys.stdin)
    else:
        try:
            f = open(fname, 'rb')
        except:
            return 'Failed to open file %s' %fname
        ret = sumfile(f)
        f.close()
    return ret

class Project_File:
    """
    Container of file information to send to postres.
   
    """
    def __init__(self):
        """
        Attributes must be the same as the columns in table project_files,
        nothing more nothing less. Attribute names must be the same as column
        names, order doesn't matter.
        
        """
        self.project= ''
        self.hostname= ''
        self.path= ''
        self.filename= ''
        self.ctime= ''
        self.mtime= ''
        self.fsize= ''
        self.md5sum= ''
        self.removed= ''
        self.description= ''
        
def stat_file(project_file, project= args.project, description= args.description):
    """
    Returns a Project_File object populated by the file "project_file"
    
    ARGS:
    project_file <str>:
        file name including *full* path
    
    project & description <str>:
        Which project does this file belongs and description
    
    RETURNS:
    Project_File object, that is the info about project_file

    EXAMPLE:
    p= stat_file('/home/berald01/Tritume/Rplots.pdf')
    print(p.filename)
    >>> Rplots.pdf

    """
    pfile= Project_File()
    pfile.project= project
    pfile.hostname= socket.gethostname()
    (pfile.path, pfile.filename)= os.path.split(project_file)
    pfile.ctime= datetime.datetime.fromtimestamp(os.path.getctime(project_file))
    pfile.mtime= datetime.datetime.fromtimestamp(os.path.getmtime(project_file))
    pfile.fsize= os.path.getsize(project_file)
    pfile.md5sum= md5sum(project_file)
    pfile.removed= False
    pfile.description= description
    return(pfile)
    

if args.files == ['-']:
    args.files= sys.stdin.readlines()
    args.files= [x.strip() for x in args.files]
    print(args.files)

project_files= []
for file in args.files:
    if os.path.isdir(file):
        if args.recursive is True:
            for root, dirnames, filenames in os.walk(file):
                for filename in filenames:
                    project_files.append(os.path.join(root, filename))
        else:
            files= os.listdir(file)
            files= [os.path.join(file, x) for x in files if os.path.isfile(os.path.join(file, x))]
            for x in files:
                project_files.append(x)
    elif os.path.isfile(file):
        project_files.append(file)
    elif os.path.islink(file):
        print("Skipping file %s (it's a link)" %(file))
        continue
    else:
        sys.exit('project_files: Item "%s" in --files list: No such file or directory' %(file) )
project_files= list(set(project_files))
project_files.sort()
project_files= [os.path.abspath(x) for x in project_files]

if args.hidden is False:
    " Remove files starting with '.' or contained in dirs starting with '.' even if the dotted dir is upstream to the current wdir"
    project_files_nothidden= []
    for f in project_files:
        fsplit= os.path.normpath(f).split(os.sep)
        dotted= [x for x in fsplit if x.startswith('.')]
        if dotted == []:
            project_files_nothidden.append(f)
        else:
            pass
else:
    project_files_nothidden= project_files

file_stats= [stat_file(x) for x in project_files_nothidden]

conn = psycopg2.connect(conn_args)
cur= conn.cursor()

for i in range(0, len(file_stats)):
    ## update_str= cur.mogrify("""UPDATE project_files SET removed= True WHERE project= %(project)s AND hostname= %(hostname)s AND path= %(path)s AND filename= %(filename)s """, vars(file_stats[i]))
    update_str= cur.mogrify("""DELETE FROM project_files WHERE project= %(project)s AND hostname= %(hostname)s AND path= %(path)s AND filename= %(filename)s """, vars(file_stats[i]))
    insert_str= cur.mogrify("""INSERT INTO project_files (project, hostname, path, filename, ctime, mtime, fsize, md5sum, removed, description) VALUES (%(project)s, %(hostname)s, %(path)s, %(filename)s, %(ctime)s, %(mtime)s, %(fsize)s, %(md5sum)s, %(removed)s, %(description)s)""", vars(file_stats[i]) )
    if args.verbose is True:
        print(update_str)
        print(insert_str)
    cur.execute(update_str)
    cur.execute(insert_str)
conn.commit()
cur.close()
conn.close()


sys.exit()

